Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Introduce LayerNorm optimization from latest Apex #277

Merged
merged 4 commits into from
Nov 2, 2023
Merged

Introduce LayerNorm optimization from latest Apex #277

merged 4 commits into from
Nov 2, 2023

Conversation

Quentin-Anthony
Copy link

Introduced in NVIDIA/apex#1715

My PR lets the user disable this LayerNorm optimization, but I suspect everyone will use it so it's on-by-default.

Not backwards-compatible with older Apex. Do you need a version check or is this ok?

@tjruwase
Copy link

tjruwase commented Nov 1, 2023

Not backwards-compatible with older Apex. Do you need a version check or is this ok?

@Quentin-Anthony, thanks for this PR. But we do need backwards-compatibility, so please add a version check.

@Quentin-Anthony
Copy link
Author

Not backwards-compatible with older Apex. Do you need a version check or is this ok?

@Quentin-Anthony, thanks for this PR. But we do need backwards-compatibility, so please add a version check.

Apex doesn't have versioning yet, so I added support to manually inspect the function and see if the memory_efficient arg exists in FusedLayerNormAffineFunction.forward, which is a bit messy but does the job.

Hopefully in the future NVIDIA/apex#1648 gets merged and we can just check apex.__version__

@tjruwase
Copy link

tjruwase commented Nov 2, 2023

Apex doesn't have versioning yet, so I added support to manually inspect the function and see if the memory_efficient arg exists in FusedLayerNormAffineFunction.forward, which is a bit messy but does the job.

Works for me. Thanks!

@tjruwase tjruwase merged commit ef13d09 into microsoft:main Nov 2, 2023
1 check passed
@RuiWang1998
Copy link

Hi,

Author of NVIDIA/apex#1715 here. Thanks for incorporate this into the repo (as the default)! This is very exciting.

Moreoever, I'm writing to let you guys know that https://github.com/Quentin-Anthony/Megatron-DeepSpeed-MS/blob/046319fecccfb8053ad3de5181e48f943ff14d27/megatron/model/fused_layer_norm.py#L96C18-L96C75 also has the same memory_efficient feature in the same pr!

@tjruwase
Copy link

@RuiWang1998, thanks for the information. @Quentin-Anthony, do you have bandwidth to handle this?

@Quentin-Anthony
Copy link
Author

@RuiWang1998, thanks for the information. @Quentin-Anthony, do you have bandwidth to handle this?

Yep I'll take care of it

@RuiWang1998 RuiWang1998 mentioned this pull request Dec 21, 2023
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@tjruwase tjruwase tjruwase approved these changes

@jeffra jeffra Awaiting requested review from jeffra

@ShadenSmith ShadenSmith Awaiting requested review from ShadenSmith

@conglongli conglongli Awaiting requested review from conglongli

@awan-10 awan-10 Awaiting requested review from awan-10 awan-10 is a code owner

@eltonzheng eltonzheng Awaiting requested review from eltonzheng

@minjiaz minjiaz Awaiting requested review from minjiaz

@RezaYazdaniAminabadi RezaYazdaniAminabadi Awaiting requested review from RezaYazdaniAminabadi

@duli2012 duli2012 Awaiting requested review from duli2012

@mrwyattii mrwyattii Awaiting requested review from mrwyattii

@yaozhewei yaozhewei Awaiting requested review from yaozhewei

@arashb arashb Awaiting requested review from arashb

@xiaoxiawu-microsoft xiaoxiawu-microsoft Awaiting requested review from xiaoxiawu-microsoft

@GuanhuaWang GuanhuaWang Awaiting requested review from GuanhuaWang GuanhuaWang is a code owner

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@Quentin-Anthony @tjruwase @RuiWang1998